Vertex AI Workbenchのマネージドノートブックでカスタムコンテナを利用する

おはこんハロチャオ～！なにもんなんじゃ？じょんすみすです。

当エントリー『クラスメソッド 機械学習チーム アドベントカレンダー 2022』は13日目となります。

みなさん、Vertex AI Workbench(以下Workbench)使ってますか？ Vertex AIの機能の一部で、Jupyterの環境を提供してくれています。

このWorkbenchにおいてマネージドノートブックは、 実際に動いてる環境をユーザが意識せずにJupyter上での開発のみに専念できます。

今回は、そんなマネージドノートブックにカスタムコンテナで実装したカーネルを追加する方法を見ていきます。

カスタムコンテナが必要な理由

やり方を見ていく前に、そもそも、なぜカスタムコンテナを使ってカーネルを追加したいのかについて考えておきます。

マネージドノートブックはそれが実行される環境のことをユーザが意識しなくて良いという利点がありますが、 これは、ユーザがその環境側に対する変更が行えないことも意味しています。 標準で用意されているカーネルに含まれる範囲で必要な処理の実装が完結する場合はいいのですが、そうでない場合に困ります。

実行する処理に応じて必要な様々なライブラリなどが別途入った環境はあらかじめすべてのパターンを網羅するのは難しいため、 マネージドノートブックは利用する環境をまとめたコンテナを別途用意して利用できる仕組みを提供しているわけです。

カスタムコンテナを作成する

まずは、Workbenchで利用可能なカスタムコンテナの作り方を見ていきましょう。

Workbenchで利用するコンテナにはJupyterで利用可能なカーネルがインストールされている必要があります。 Pythonであれば ipykernel 、　Rであれば IRkernel といったものになります。

言語そのものを動かす環境はもちろん必要ですが、必須となる要件はそれだけです。 これは、通常のJuypter環境において、カーネルを追加する際と同様と考えていただければ問題ありません。

WorkbenchにはDeep Learning Containerを使ったカーネルが標準で含まれています。

• Choose a container image  |  Deep Learning Containers  |  Google Cloud

PythonやRであれば、これを利用することもできます。 例えば、MeCabを追加したコンテナを作成する場合は以下のように利用できます。

FROM gcr.io/deeplearning-platform-release/base-cpu:latest

RUN apt update && \
    apt upgrade -y && \
    apt install -y --no-install-recommends \
        mecab \
        libmecab-dev \
        mecab-ipadic \
        mecab-ipadic-utf8
RUN pip install mecab-python3

ENV MECABRC /etc/mecabrc

このコンテナをContainer RegistryやArtifact Registryに登録することで、Workbench内のPythonカーネルにMeCabが追加された新しいカーネルを追加する準備が整います。

Workbenchから利用できるようにする

コンテナの準備が整ったら、これをWorkbenchから利用できるように設定します。

まずはマネージドノートブックを作成します。

新規作成画面で「詳細設定」を展開し、「カスタムDockerイメージを指定する」にチェックを入れます。 コンテナなパスを指定する項目が出てくるので、先ほど登録したコンテナを指定します。

この設定をしたのちノートブックを作成すると、ノートブック一覧に以下のようなイメージを追加したものが表示され作成が開始されます。

作成が完了し、「JUPYTERLABを開く」からノートブック環境を開くと右側にカスタムコンテナを取り込み中であるメッセージが表示されます。

カスタムコンテナの取り込みが完了すると、カーネルとして利用可能になります。

おわりに

今回は、Vertex AI Workbenchでカスタムコンテナを使ったカーネルの追加方法を見ていきました。

既存の環境だけだと足りない部分がある場合や、Workbenchのローカルで環境を合わせたいなどの場合にご活用ください。